#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
QTorch量化交易框架 - 多资产回测示例
支持股票、期货、期权等多种投资标的同时回测
支持GPU加速计算
"""

import logging
import argparse
import os
import yaml
import pandas as pd
import time
from qtorch.core.qengine import QEngine
from qtorch.strategy.simple_sma_strategy import SimpleSMAStrategy
from qtorch.strategy.dual_moving_average import DualMovingAverageStrategy
from qtorch.utils.device_utils import get_device, is_gpu_available

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='QTorch量化交易回测系统 - 多资产回测')
    parser.add_argument('--stocks', type=str, default='000001,600000', help='股票代码，逗号分隔')
    parser.add_argument('--futures', type=str, default='IF888.CFX', help='期货代码，逗号分隔')
    parser.add_argument('--options', type=str, default='', help='期权代码，逗号分隔')
    parser.add_argument('--start', type=str, default='2020-01-01', help='回测开始日期')
    parser.add_argument('--end', type=str, default='2022-12-31', help='回测结束日期')
    parser.add_argument('--fast', type=int, default=5, help='快线周期')
    parser.add_argument('--slow', type=int, default=20, help='慢线周期')
    parser.add_argument('--cash', type=float, default=1000000, help='初始资金')
    parser.add_argument('--output', type=str, default='./output/multi_asset', help='结果输出路径')
    
    # GPU加速相关参数
    parser.add_argument('--force-cpu', action='store_true', help='强制使用CPU，即使有GPU可用')
    parser.add_argument('--gpu-stats', action='store_true', help='显示GPU性能统计')
    parser.add_argument('--mixed-precision', action='store_true', help='使用混合精度计算')
    return parser.parse_args()

def show_gpu_info():
    """显示GPU信息"""
    if is_gpu_available():
        device = get_device()
        logging.info(f"已检测到可用GPU: {device}")
        
        # 如果是CUDA设备，显示更多信息
        if hasattr(device, 'type') and device.type == 'cuda':
            import torch
            for i in range(torch.cuda.device_count()):
                props = torch.cuda.get_device_properties(i)
                logging.info(f"  CUDA:{i} - {props.name}, 内存: {props.total_memory / 1024**3:.1f}GB")
    else:
        logging.info("未检测到GPU，将使用CPU进行计算")

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_args()
    
    # 显示GPU信息
    show_gpu_info()
    
    start_time = time.time()
    
    try:
        # 创建回测引擎
        engine = QEngine(force_cpu=args.force_cpu)
        
        # 配置GPU加速
        if not args.force_cpu and is_gpu_available():
            engine.set_mixed_precision(args.mixed_precision)
            logging.info("已启用GPU加速")
        
        # 添加股票数据
        if args.stocks:
            for symbol in args.stocks.split(','):
                if symbol:
                    engine.add_data(
                        symbol=symbol,
                        start_date=args.start,
                        end_date=args.end,
                        asset_type='stock'
                    )
                    logging.info(f"添加股票: {symbol}")
        
        # 添加期货数据
        if args.futures:
            for symbol in args.futures.split(','):
                if symbol:
                    engine.add_data(
                        symbol=symbol,
                        start_date=args.start,
                        end_date=args.end,
                        asset_type='futures'
                    )
                    logging.info(f"添加期货: {symbol}")
        
        # 添加期权数据
        if args.options:
            for symbol in args.options.split(','):
                if symbol:
                    engine.add_data(
                        symbol=symbol,
                        start_date=args.start,
                        end_date=args.end,
                        asset_type='option'
                    )
                    logging.info(f"添加期权: {symbol}")
        
        # 设置回测参数
        engine.set_cash(args.cash)
        
        # 设置不同资产类型的交易佣金
        engine.set_commission(0.0003, asset_type='stock')    # 股票佣金
        engine.set_commission(0.00005, asset_type='futures') # 期货佣金
        engine.set_commission(0.0005, asset_type='option')   # 期权佣金
        
        # 设置不同资产类型的保证金比例
        engine.set_margin(0.12, asset_type='stock')     # 股票保证金
        engine.set_margin(0.08, asset_type='futures')   # 期货保证金
        engine.set_margin(1.0, asset_type='option')     # 期权保证金(全额)
        
        # 设置不同资产类型的合约乘数
        engine.set_contract_multiplier(1.0, asset_type='stock')      # 股票乘数
        engine.set_contract_multiplier(300.0, asset_type='futures')  # 期货乘数(如IF)
        engine.set_contract_multiplier(10000.0, asset_type='option') # 期权乘数
        
        # 设置是否允许做空
        engine.allow_short(True)
        
        # 添加策略
        engine.add_strategy(
            DualMovingAverageStrategy,
            fast_period=args.fast,
            slow_period=args.slow
        )
        
        # 运行回测
        logging.info(f"开始多资产回测，从 {args.start} 到 {args.end}")
        result = engine.run()
        
        # 打印结果摘要
        result.summary()
        
        # 显示GPU性能统计（如果启用）
        if args.gpu_stats:
            result.print_performance_stats()
        
        # 计算总运行时间
        total_time = time.time() - start_time
        logging.info(f"多资产回测完成，耗时: {total_time:.2f}秒")
        
        # 确保输出目录存在
        os.makedirs(args.output, exist_ok=True)
        
        # 生成回测报告
        result.plot(output_dir=args.output, show_gpu_stats=args.gpu_stats)
        
        # 保存回测数据和指标
        metrics_path = os.path.join(args.output, "metrics.yaml")
        with open(metrics_path, 'w', encoding='utf-8') as f:
            yaml.dump(result.metrics, f, default_flow_style=False, allow_unicode=True)
            
        result_path = os.path.join(args.output, "result.csv")
        result.data.to_csv(result_path)
        
        logging.info(f"多资产回测完成，报告已保存至 {args.output}")
        
    except Exception as e:
        logging.error(f"多资产回测执行失败: {str(e)}", exc_info=True)
        raise

if __name__ == '__main__':
    main()